package org.example;


import java.io.IOException;
import java.net.InetSocketAddress;
import java.net.SocketAddress;
import java.nio.ByteBuffer;
import java.nio.channels.AsynchronousCloseException;
import java.nio.channels.AsynchronousSocketChannel;
import java.nio.channels.CompletionHandler;
import java.util.Scanner;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.Future;
import java.util.concurrent.TimeUnit;

public class Main {

    private static final String dstHost = "192.168.1.10";
    private static final int dstPort = 16088;

    private static final String srcHost = "localhost";
    private static final int srcPort = 3389;


    public static void main(String[] args) throws IOException {

        System.out.println("dst:"+dstHost+":"+dstPort);
        System.out.println("src:"+srcHost+":"+srcPort);

        //使用aio
        final AsynchronousSocketChannel client = AsynchronousSocketChannel.open();

        client.connect(new InetSocketAddress(dstHost, dstPort), null, new CompletionHandler<Void, Void>() {
            public void completed(Void result, Void attachment) {
                //连接成功
                byte[] bt = new byte[]{0};
                final ByteBuffer buffer = ByteBuffer.wrap(bt);
                client.write(buffer, null, new CompletionHandler<Integer, Void>() {
                    public void completed(Integer result, Void attachment) {

                        //读取数据
                        final ByteBuffer buffer = ByteBuffer.allocate(1);
                        client.read(buffer, null, new CompletionHandler<Integer, Void>() {
                            public void completed(Integer result, Void attachment) {
                                buffer.flip();

                                if (buffer.get() == 1) {
                                    //发起新的连
                                    try {
                                        createNewClient();
                                    } catch (IOException e) {
                                        throw new RuntimeException(e);
                                    }
                                }
                                buffer.clear();
                                // 这里再次调用读取操作，实现循环读取
                                client.read(buffer, null, this);
                            }

                            public void failed(Throwable exc, Void attachment) {
                                exc.printStackTrace();
                            }
                        });


                    }

                    public void failed(Throwable exc, Void attachment) {
                        exc.printStackTrace();
                    }
                });


            }

            public void failed(Throwable exc, Void attachment) {
                exc.printStackTrace();
            }
        });
        Scanner in = new Scanner(System.in);
        in.nextLine();

    }

    private static void createNewClient() throws IOException {

        final AsynchronousSocketChannel dstClient = AsynchronousSocketChannel.open();
        dstClient.connect(new InetSocketAddress(dstHost, dstPort), null, new CompletionHandler<Void, Void>() {
            public void completed(Void result, Void attachment) {

                //尝试连接本地服务
                final AsynchronousSocketChannel srcClient;
                try {
                    srcClient = AsynchronousSocketChannel.open();
                    srcClient.connect(new InetSocketAddress(srcHost, srcPort), null, new CompletionHandler<Void, Void>() {
                        public void completed(Void result, Void attachment) {

                            byte[] bt = new byte[]{1};
                            final ByteBuffer buffer = ByteBuffer.wrap(bt);
                            Future<Integer> write = dstClient.write(buffer);
                            try {
                                write.get();
                                //交换数据
                                exchangeData(srcClient, dstClient);
                                exchangeData(dstClient, srcClient);
                            } catch (Exception e) {
                                closeChannels(srcClient, dstClient);
                            }


                        }

                        public void failed(Throwable exc, Void attachment) {
                            exc.printStackTrace();
                            //失败
                            byte[] bt = new byte[]{2};
                            final ByteBuffer buffer = ByteBuffer.wrap(bt);
                            dstClient.write(buffer);
                        }
                    });

                } catch (IOException e) {
                    e.printStackTrace();
                    //失败
                    byte[] bt = new byte[]{2};
                    final ByteBuffer buffer = ByteBuffer.wrap(bt);
                    dstClient.write(buffer);
                }

            }

            public void failed(Throwable exc, Void attachment) {
                exc.printStackTrace();
            }
        });
    }

    private static void exchangeData(AsynchronousSocketChannel ch1, AsynchronousSocketChannel ch2) {
        try {
            final ByteBuffer buffer = ByteBuffer.allocate(1024);

            ch1.read(buffer, null, new CompletionHandler<Integer, CompletableFuture<Integer>>() {

                public void completed(Integer result, CompletableFuture<Integer> readAtt) {

                    CompletableFuture<Integer> future = new CompletableFuture<>();

                    if (result == -1 || buffer.position() == 0) {
                        // 处理连接关闭的情况或者没有数据可读的情况

                        try {
                            readAtt.get(3,TimeUnit.SECONDS);
                        } catch (Exception e) {
                            e.printStackTrace();
                        }

                        closeChannels(ch1, ch2);
                        return;
                    }

                    buffer.flip();

                    CompletionHandler readHandler = this;

                    ch2.write(buffer, future, new CompletionHandler<Integer, CompletableFuture<Integer>>() {
                        @Override
                        public void completed(Integer result, CompletableFuture<Integer> writeAtt) {

                            if (buffer.hasRemaining()) {
                                // 如果未完全写入，则继续写入
                                ch2.write(buffer, writeAtt, this);

                            } else {
                                writeAtt.complete(1);
                                // 清空buffer并继续读取
                                buffer.clear();
                                if(ch1.isOpen()){
                                    ch1.read(buffer, writeAtt, readHandler);
                                }
                            }

                        }

                        @Override
                        public void failed(Throwable exc, CompletableFuture<Integer> attachment) {
                            if(!(exc instanceof AsynchronousCloseException)){
                                exc.printStackTrace();
                            }
                            closeChannels(ch1, ch2);
                        }
                    });

                }

                public void failed(Throwable exc, CompletableFuture<Integer>  attachment) {
                    if(!(exc instanceof AsynchronousCloseException)){
                        exc.printStackTrace();
                    }
                    closeChannels(ch1, ch2);
                }
            });

        } catch (Exception ex) {
            ex.printStackTrace();
            closeChannels(ch1, ch2);
        }

    }

    private static void closeChannels(AsynchronousSocketChannel ch1, AsynchronousSocketChannel ch2) {
        if (ch1 != null && ch1.isOpen()) {
            try {
                ch1.close();
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
        if (ch2 != null && ch2.isOpen()) {
            try {
                ch2.close();
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
    }
}